import requests
import pandas as pd
import matplotlib.pyplot as plt


class LightCurveFetcher:
    def __init__(self, base_url):
        """
        初始化光变曲线抓取器
        :param base_url: API 接口地址
        """
        self.base_url = base_url

    def fetch(self, ra, dec, radius=1.0, band='all', start_date=None, download=False):
        """
        获取光变曲线数据
        :param ra: 赤经 (float 或 "hh:mm:ss")
        :param dec: 赤纬 (float 或 "dd:mm:ss")
        :param radius: 查询半径，单位为角秒
        :param band: 波段 ('all', 'r', 'g', etc.)
        :param start_date: 起始日期 (YYYY-MM-DD 或 YYYYMMDD)
        :param download: 是否以 CSV 格式下载（True 返回文件名，否则返回 DataFrame）
        """
        params = {
            'ra': str(ra),
            'dec': str(dec),
            'radius': radius,
            'band': band,
        }
        if start_date:
            params['date_'] = start_date
        if download:
            params['download'] = 'true'

        response = requests.get(self.base_url, params=params)

        if response.status_code != 200:
            raise RuntimeError(f"请求失败: {response.status_code}\n{response.text}")

        if download:
            filename = "light_curve_data.csv"
            with open(filename, "w", encoding="utf-8") as f:
                f.write(response.text)
            return filename
        else:
            return pd.DataFrame(response.json())


def plot_light_curve(df, title="Light Curve"):
    """
    单波段光变曲线（Magnitude vs MJD）
    """
    if df.empty:
        print("⚠️ 数据为空，无法绘图。")
        return

    plt.figure(figsize=(10, 6))
    plt.errorbar(
        df['mjd'], df['mag'], 
        yerr=df['magerr_auto_s'], fmt='o',
        ecolor='gray', capsize=2, label='Photometry'
    )
    plt.gca().invert_yaxis()
    plt.xlabel("MJD")
    plt.ylabel("Magnitude")
    plt.title(title)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()


def plot_multiband(df,title=None):
    """
    多波段光变曲线（不同颜色区分 band）
    """
    if df.empty:
        print("⚠️ 数据为空，无法绘图。")
        return

    plt.figure(figsize=(10, 6))
    for band in df['band'].unique():
        sub = df[df['band'] == band]
        plt.errorbar(
            sub['mjd'], sub['mag'], yerr=sub['magerr_auto_s'],
            fmt='o', label=f'Band: {band}', capsize=2
        )
    plt.gca().invert_yaxis()
    plt.xlabel("MJD")
    plt.ylabel("Magnitude")
    if title:
        plt.title(title)
    else:
        plt.title("Multi-band Light Curve")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    # ✅ 配置参数
    BASE_URL = "http://192.168.16.70:8082/api/query_light_curve_xp_version_V20250303_KM"
    RA = '2:28:16.5000'
    DEC = '+19:11:39.900'
    RADIUS = 1.5      # arcsec
    BAND = "all"      # 'r', 'g', 'all'
    START_DATE = None # or "2025-07-01"
    DOWNLOAD = False  # 是否导出为 CSV 文件

    fetcher = LightCurveFetcher(BASE_URL)
    df = fetcher.fetch(RA, DEC, radius=RADIUS, band=BAND, start_date=START_DATE, download=DOWNLOAD)

    if DOWNLOAD:
        print(f"✅ 光变曲线 CSV 文件已保存到: {df}")
    else:
        print(f"✅ 共获取 {len(df)} 条记录")
        if BAND == "all":
            plot_multiband(df)
        else:
            plot_light_curve(df, title=f"Light Curve @ RA={RA}, Dec={DEC}")
